package com.btl.service.ai.support.rag.embedding.strategy.ingestor;

import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingModelTypeEnum;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingStoreIngestorStrategy;
import com.btl.service.ai.support.rag.store.vector.collections.VectorStoreQdrantFactory;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.IngestionResult;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;

import java.util.List;


@Slf4j
@Service
public class AllMiniLmL6V2EmbeddingStoreIngestorStrategy implements EmbeddingStoreIngestorStrategy {

    private final EmbeddingModel allMiniLmL6V2EmbeddingModel;
    private final VectorStoreQdrantFactory vectorStoreQdrantFactory;

    public AllMiniLmL6V2EmbeddingStoreIngestorStrategy(@Qualifier("allMiniLmL6V2EmbeddingModel") EmbeddingModel allMiniLmL6V2EmbeddingModel, VectorStoreQdrantFactory vectorStoreQdrantFactory) {
        this.allMiniLmL6V2EmbeddingModel = allMiniLmL6V2EmbeddingModel;
        this.vectorStoreQdrantFactory = vectorStoreQdrantFactory;
    }

    @Override
    public EmbeddingModel embeddingModel() {
        return allMiniLmL6V2EmbeddingModel;
    }

    @Override
    public EmbeddingStore<TextSegment> embeddingStore(String collectionName) {
        return vectorStoreQdrantFactory.getQdrantClient(collectionName.concat("_")
                .concat(String.valueOf(allMiniLmL6V2EmbeddingModel.dimension())));
    }

    @Override
    public IngestionResult embed(String collectionName, DocumentSplitter documentSplitter, List<Document> documents) {
        EmbeddingStore<TextSegment> embeddingStore = this.embeddingStore(collectionName);
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .embeddingModel(allMiniLmL6V2EmbeddingModel)
                .embeddingStore(embeddingStore)
                .documentSplitter(documentSplitter)
                .build();
        return ingestor.ingest(documents);
    }

    @Override
    public IngestionResult embed(String collectionName, List<Document> documents) {
        return this.embed(collectionName, null, documents);
    }

    @Override
    public EmbeddingModelTypeEnum modelType() {
        return EmbeddingModelTypeEnum.LOCAL_ALL_MINI_LM_L6_V2_EMBEDDING_MODEL;
    }
}
